import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain.memory import ConversationEntityMemory
from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE
from langchain.chains import ConversationChain

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

llm = ChatZhipuAI(model=model,
                  temperature=0.9,              
    )

memory=ConversationEntityMemory(llm=llm)
conversation = ConversationChain(
    llm=llm, 
    #verbose=True, 
    prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,
    memory=memory
)

conference_recording = """
你是一个会议的主持人，当前出席会议的嘉宾简要信息如下。你需要根据听众提问回复相关问题。
张三：人工智能专家，人工智能学院院长、教授；
李四：大数据专家，信息学院教授；
王五：软件工程专家，信息学院客座教授，智朗科技总工程师
"""

ret = conversation.predict(input=conference_recording)
print(f"\nRound-1:\n{memory.entity_cache}\n{memory.entity_store}\n\n【主持人】{ret}")
ret = conversation.predict(input="嘉宾中哪几位是教授？")
print(f"\nRound-2:\n{memory.entity_cache}\n{memory.entity_store}\n\n【主持人】{ret}")
ret = conversation.predict(input="我想了解软件在工程化方面的问题，哪位嘉宾最权威？")
print(f"\nRound-3:\n{memory.entity_cache}\n{memory.entity_store}\n\n【主持人】{ret}")